package org.demos.springaiargdemo1.service;

import org.apache.pdfbox.Loader;
import org.apache.pdfbox.io.IOUtils;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.transformer.splitter.TextSplitter;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Mono;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.core.io.Resource;
import org.springframework.http.codec.multipart.FilePart;
import reactor.core.publisher.Mono;

import java.io.*;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Collections;
import java.util.List;

@Service
public class DocumentService {
    
    private static final Logger log = LoggerFactory.getLogger(DocumentService.class);
    private final VectorStore vectorStore;

    // 文档切片设置
    private static final int CHUNK_SIZE = 1000;
    private static final int CHUNK_OVERLAP = 200;

    public DocumentService(VectorStore vectorStore ) {
        this.vectorStore = vectorStore;
    }

    /**
     * 处理上传的文档
     */
    public Mono<String> processDocument(FilePart filePart) {
        Mono<Resource> resourceMono = filePartToResource(filePart);
        Mono<List<Document>> documentMono = resourceMono.flatMap(resource -> {
            String filename = filePart.filename().toLowerCase();
                    if (filename.endsWith(".pdf")) {
                        return readPdf(resource);
                    }else if (filename.endsWith(".md")) {
                        return readMarkdown(resource);
                    } else {
                        throw new IllegalArgumentException("不支持的文件类型: " + filename);
                    }

        });
        return documentMono.map(documents -> {
            // 文档切片
            TextSplitter textSplitter = new TokenTextSplitter(CHUNK_SIZE, CHUNK_OVERLAP, 5, 10000, true);
            List<Document> splitDocuments = textSplitter.apply(documents);
            log.info("文档切片成功，生成 {} 个文档块", splitDocuments.size());

            // 存储到向量数据库
            vectorStore.add(splitDocuments);
            log.info("文档块成功存储到向量数据库");
            return "Successfully processed document: " + filePart.filename();
        });
    }



    /**
     * 获取文件扩展名
     */
    private String getFileExtension(String fileName) {
        int lastDotIndex = fileName.lastIndexOf('.');
        return lastDotIndex > 0 ? fileName.substring(lastDotIndex) : "";
    }


    public Mono<List<Document>> readPdf(Resource resource) {
        return Mono.fromCallable(() -> {
            PdfDocumentReaderConfig config = PdfDocumentReaderConfig.builder()
                    .withPageExtractedTextFormatter(new ExtractedTextFormatter.Builder()
                            .withNumberOfTopTextLinesToDelete(0)
                            .build())
                    .build();
            PagePdfDocumentReader reader = new PagePdfDocumentReader(resource, config);
            return reader.get();
        });
    }

    public Mono<List<Document>> readMarkdown(Resource resource) {
        return Mono.fromCallable(() -> {
            MarkdownDocumentReader reader = new MarkdownDocumentReader(resource, MarkdownDocumentReaderConfig.defaultConfig());
            return reader.get();
        });
    }

    public  Mono<Resource> filePartToResource(FilePart filePart) {
        return filePart.content()
                .collectList()
                .map(dataBuffers -> {
                    try (InputStream inputStream = dataBuffersToInputStream(dataBuffers)) {
                        byte[] bytes = IOUtils.toByteArray(inputStream);
                        return new ByteArrayResource(bytes) {
                            @Override
                            public String getFilename() {
                                return filePart.filename();
                            }
                        };
                    } catch (IOException e) {
                        throw new RuntimeException("Error converting file to resource", e);
                    } finally {
                        dataBuffers.forEach(DataBufferUtils::release);
                    }
                });
    }

    private static InputStream dataBuffersToInputStream(List<DataBuffer> dataBuffers) {
        return new SequenceInputStream(Collections.enumeration(
                dataBuffers.stream()
                        .map(dataBuffer -> new ByteArrayInputStream(
                                dataBuffer.asByteBuffer().array(),
                                dataBuffer.readPosition(),
                                dataBuffer.readableByteCount()
                        ))
                        .toList()
        ));
    }
}
